#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
行业指标数据提取工具
通过命令行参数根据行业关键词和年份提取指标数据
"""

import pandas as pd
import json
import argparse
import sys
from pathlib import Path

def load_data_and_dictionary(data_file, dict_file):
    """加载数据文件和字典文件"""
    try:
        # 读取行业指标数据
        if data_file.endswith('.xlsx'):
            if Path(data_file).exists():
                df_data = pd.read_excel(data_file, sheet_name='整理后数据')
            else:
                df_data = pd.read_excel(data_file)
        else:
            print(f"❌ 不支持的文件格式: {data_file}")
            return None, None
        
        # 读取指标字典
        df_dict = None
        if Path(dict_file).exists():
            df_dict = pd.read_excel(dict_file)
        else:
            print(f"⚠️  警告: 字典文件不存在 {dict_file}")
            print("将使用基础分类字典")
            # 使用内置的分类字典
            classification = {
                'EBITDA率(%)': '盈利能力分析',
                '净资产收益率(%)': '盈利能力分析',
                '总资产收益率(%)': '盈利能力分析',
                '销售营业利润率(%)': '盈利能力分析',
                '国有资本回报率(%)': '盈利能力分析',
                '经济增加值率(%)': '盈利能力分析',
                '资本保值增值率(%)': '盈利能力分析',
                '存货周转率(次)': '经营周转能力分析',
                '应收账款周转率(次)': '经营周转能力分析',
                '总资产周转率(次)': '经营周转能力分析',
                '流动资产周转率(次)': '经营周转能力分析',
                '带息负债比率(%)': '偿债能力分析',
                '资产负债率(%)': '偿债能力分析',
                '速动比率': '偿债能力分析',
                '现金流动负债比率(%)': '偿债能力分析',
                '已获利息倍数(倍)': '偿债能力分析',
                '盈余现金保障倍数(倍)': '偿债能力分析',
                '营业现金比率(%)': '偿债能力分析',
                '利润总额增长率(%)': '成长能力分析',
                '营业总收入增长率(%)': '成长能力分析',
                '技术投入比率(%)': '资本性支出分析',
                '研发经费投入强度(%)': '资本性支出分析',
                '百元收入支付的成本费用(元)': '成本控制分析',
                '两金占流动资产比重(%)': '成本控制分析',
                '全员劳动生产率(万元/人)': '人力资源效率分析'
            }
            df_dict = pd.DataFrame([
                {'指标名称': k, '指标大类': v} for k, v in classification.items()
            ])
        
        return df_data, df_dict
        
    except Exception as e:
        print(f"❌ 加载数据文件时出错: {e}")
        return None, None

def filter_data(df_data, df_dict, industry_keyword, year):
    """根据条件筛选数据"""
    
    # 筛选行业（关键词匹配）
    if industry_keyword:
        mask = df_data['行业名称'].str.contains(industry_keyword, case=False, na=False)
        df_filtered = df_data[mask].copy()
        
        if len(df_filtered) == 0:
            print(f"❌ 没有找到包含关键词 '{industry_keyword}' 的行业")
            # 显示可用的行业列表
            print("可用的行业:")
            for industry in sorted(df_data['行业名称'].unique()):
                print(f"  - {industry}")
            return None
    else:
        df_filtered = df_data.copy()
    
    # 筛选年份
    if year:
        df_filtered = df_filtered[df_filtered['年份'] == year]
        
        if len(df_filtered) == 0:
            print(f"❌ 没有找到 {year} 年的数据")
            available_years = sorted(df_data['年份'].unique())
            print(f"可用年份: {available_years}")
            return None
    
    # 关联指标字典获取指标大类
    if df_dict is not None:
        # 创建指标名称到指标大类的映射
        indicator_mapping = dict(zip(df_dict['指标名称'], df_dict['指标大类']))
        df_filtered['指标大类'] = df_filtered['指标名称'].map(indicator_mapping)
        
        # 如果有指标没有找到对应的大类，设为"未分类"
        df_filtered['指标大类'] = df_filtered['指标大类'].fillna('未分类')
    else:
        df_filtered['指标大类'] = '未知'
    
    return df_filtered

def convert_to_json(df_filtered, output_format='pretty'):
    """将筛选后的数据转换为JSON格式"""
    
    if df_filtered is None or len(df_filtered) == 0:
        return None
    
    # 重新排列列顺序
    columns_order = ['行业名称', '年份', '指标大类', '指标名称', '优秀值', '良好值', '平均值', '较低值', '较差值']
    df_output = df_filtered[columns_order].copy()
    
    # 转换为字典列表
    data_list = df_output.to_dict('records')
    
    # 处理NaN值
    for record in data_list:
        for key, value in record.items():
            if pd.isna(value):
                record[key] = None
    
    # 创建结果字典
    result = {
        "query_info": {
            "total_records": len(data_list),
            "industries": sorted(df_filtered['行业名称'].unique().tolist()),
            "years": sorted(df_filtered['年份'].unique().tolist()),
            "indicators": sorted(df_filtered['指标名称'].unique().tolist())
        },
        "data": data_list
    }
    
    if output_format == 'pretty':
        return json.dumps(result, ensure_ascii=False, indent=2)
    else:
        return json.dumps(result, ensure_ascii=False)

def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description='行业指标数据提取工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  python extract_industry_data.py --industry 白酒 --year 2023
  python extract_industry_data.py --industry 制造 --year 2022 --output result.json
  python extract_industry_data.py --industry 零售 --data custom_data.xlsx
  python extract_industry_data.py --list-industries
  python extract_industry_data.py --list-years
        """
    )
    
    parser.add_argument('--industry', '-i', 
                       help='行业关键词（支持模糊匹配）')
    parser.add_argument('--year', '-y', type=int,
                       help='指定年份')
    parser.add_argument('--data', '-d', default='行业均值/整理后的行业指标数据.xlsx',
                       help='数据文件路径（默认: 行业均值/整理后的行业指标数据.xlsx）')
    parser.add_argument('--dict', default='行业均值/行业指标字典表.xlsx',
                       help='字典文件路径（默认: 行业均值/行业指标字典表.xlsx）')
    parser.add_argument('--output', '-o',
                       help='输出文件路径（不指定则输出到控制台）')
    parser.add_argument('--format', choices=['pretty', 'compact'], default='pretty',
                       help='JSON输出格式（默认: pretty）')
    parser.add_argument('--list-industries', action='store_true',
                       help='列出所有可用的行业')
    parser.add_argument('--list-years', action='store_true',
                       help='列出所有可用的年份')
    
    args = parser.parse_args()
    
    # 加载数据
    print("📊 正在加载数据...")
    df_data, df_dict = load_data_and_dictionary(args.data, args.dict)
    
    if df_data is None:
        sys.exit(1)
    
    print(f"✅ 数据加载完成: {len(df_data)} 条记录")
    
    # 处理列表请求
    if args.list_industries:
        print("\n🏢 可用行业列表:")
        for industry in sorted(df_data['行业名称'].unique()):
            print(f"  - {industry}")
        return
    
    if args.list_years:
        print("\n📅 可用年份列表:")
        years = sorted(df_data['年份'].unique())
        print(f"  {years}")
        return
    
    # 验证必需参数
    if not args.industry and not args.year:
        print("❌ 请至少指定行业关键词或年份")
        print("使用 --help 查看帮助信息")
        sys.exit(1)
    
    # 筛选数据
    print(f"🔍 筛选条件: 行业='{args.industry}', 年份={args.year}")
    df_filtered = filter_data(df_data, df_dict, args.industry, args.year)
    
    if df_filtered is None or len(df_filtered) == 0:
        print("❌ 没有找到符合条件的数据")
        sys.exit(1)
    
    print(f"✅ 找到 {len(df_filtered)} 条符合条件的记录")
    
    # 转换为JSON
    json_result = convert_to_json(df_filtered, args.format)
    
    if json_result is None:
        print("❌ JSON转换失败")
        sys.exit(1)
    
    # 输出结果
    if args.output:
        try:
            with open(args.output, 'w', encoding='utf-8') as f:
                f.write(json_result)
            print(f"✅ 结果已保存到: {args.output}")
        except Exception as e:
            print(f"❌ 保存文件失败: {e}")
            sys.exit(1)
    else:
        print("\n📋 查询结果:")
        print(json_result)

if __name__ == "__main__":
    main()
